#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
import os
from enum import Enum

LOG_PRIVILEGE = 0o640
LOG_DIR_PRIVILEGE = 0o750
DEFAULT_LOG_SUB_DIR = "training_toolkit_default_log_dir"

WORK_DIR = os.path.dirname(os.path.dirname(__file__))
LOG_PATH = "/var/log/mindx-dl/training_toolkit"
ASCEND_DEFAULT_INSTALL_PATH = "/usr/local/Ascend"
ASCEND_TOOLKIT_DEFAULT_INSTALL_PATH = f"{ASCEND_DEFAULT_INSTALL_PATH}/ascend-toolkit"
ASCEND_NNAE_DEFAULT_INSTALL_PATH = f"{ASCEND_DEFAULT_INSTALL_PATH}/nnae"

ASCEND_TOOLKIT_DEFAULT_SET_ENV_PATH = f"{ASCEND_TOOLKIT_DEFAULT_INSTALL_PATH}/set_env.sh"
ASCEND_NNAE_DEFAULT_SET_ENV_PATH = f"{ASCEND_NNAE_DEFAULT_INSTALL_PATH}/set_env.sh"
ASCEND_TF_PLUGIN_DEFAULT_SET_ENV_PATH = f"{ASCEND_DEFAULT_INSTALL_PATH}/tfplugin/set_env.sh"

# environment varibale
ASCEND_AICPU_PATH_ENV_KEY = "ASCEND_AICPU_PATH"
LD_LIBRARY_PATH_ENV_KEY = "LD_LIBRARY_PATH"
PATH_ENV_KEY = "PATH"
OPTION_EXEC_EXTERN_PLUGIN_PATH_ENV_KEY = "OPTION_EXEC_EXTERN_PLUGIN_PATH"
HOST_UP_ENV_KEY = "XDL_IP"
DEVICE_ID_ENV_KEY = "DEVICE_ID"
ASCEND_DEVICE_ID_ENV_KEY = "ASCEND_DEVICE_ID"
RANK_ID_ENV_KEY = "RANK_ID"
DEVICE_INDEX_ENV_KEY = "DEVICE_INDEX"
RANK_SIZE_ENV_KEY = "RANK_SIZE"
NODE_NUM_ENV_KEY = "NODE_NUM"
JOB_ID_ENV_KEY = "JOB_ID"
TASK_QUEUE_ENABLE_ENV_KEY = "TASK_QUEUE_ENABLE"
PYTORCH_MASTER_IP_ENV_KEY = "MASTER_ADDR"
PYTORCH_MASTER_PORT_ENV_KEY = "MASTER_PORT"
PYTORCH_MASTER_PORT_ENV_VALUE = "29501"
PYTORCH_RANK_ENV_KEY = "RANK"
PYTORCH_WORLD_SIZE_ENV_KEY = "WORLD_SIZE"
HCCL_WHITELIST_DISABLE_ENV_KEY = "HCCL_WHITELIST_DISABLE"
HCCL_IF_IP_ENV_KEY = "HCCL_IF_IP"
COMBINED_ENBALE_ENV_KEY = "COMBINED_ENBALE"
DEVICE_RANK_LIST_ENV_KEY = "DEVICE_RANK_LIST"
TERMINATION_GRACE_PERIOD_SECONDS_ENV_KEY = "TERMINATION_GRACE_PERIOD_SECONDS"
TERMINATION_GRACE_PERIOD_SECONDS_ENV_VAL = "600"
PYTHONUNBUFFERED_EN_KEY = "PYTHONUNBUFFERED"
PYTHONPATH_ENV_KEY = "PYTHONPATH"
ASCEND_SLOG_PRINT_TO_STDOUT_ENV_KEY = "ASCEND_SLOG_PRINT_TO_STDOUT"
HCCL_CONNECT_TIMEOUT_ENV_KEY = "HCCL_CONNECT_TIMEOUT"
HCCL_EXEC_TIMEOUT_ENV_KEY = "HCCL_EXEC_TIMEOUT"
HCCL_CONNECT_TIMEOUT_ENV_VALUE = "120"
HCCL_EXEC_TIMEOUT_ENV_VALUE = "600"

# required env for tensorflow without ranktable
CM_CHIEF_IP_ENV_KEY = "CM_CHIEF_IP"
CM_CHIEF_PORT_ENV_KEY = "CM_CHIEF_PORT"
CM_CHIEF_DEVICE_ENV_KEY = "CM_CHIEF_DEVICE"
CM_WORKER_IP_ENV_KEY = "CM_WORKER_IP"
CM_WORKER_SIZE_ENV_KEY = "CM_WORKER_SIZE"

CHECK_INTERVAL = 5
MAX_RANK_SIZE = 4096
MAX_CHIP_ONE_NODE = 8
MIN_GRACE_TIME = 30
MAX_GRACE_TIME = 1800


class PlatformType(Enum):
    Pytorch: str = "pt"
    MindSpore: str = "ms"
    TensorFlow: str = "tf"
    TensorFlow2: str = "tf2"
